#pragma once

#include <vector>
#include <cstdlib>
#include <cmath>
#include <stdexcept>
#include <iostream>

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<class TYPE>
class LeastSquaresRegression {
public:
    LeastSquaresRegression() = default;

    ~LeastSquaresRegression() = default;

    bool fitIt(const std::vector<TYPE> &x, const std::vector<TYPE> &y, std::vector<TYPE> &coeffs);
};

////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
template<class TYPE>
bool
LeastSquaresRegression<TYPE>::fitIt(const std::vector<TYPE> &x, const std::vector<TYPE> &y, std::vector<TYPE> &coeffs) {
    // The size of x and y must be the same
    if (x.size() != y.size()) {
        throw std::runtime_error("The size of x & y arrays are different");
    }

    // The size of x and y cannot be 0
    if (x.size() == 0 || y.size() == 0) {
        throw std::runtime_error("The size of x or y arrays is 0");
    }

    size_t N = x.size();

    // Sum of x and y
    TYPE sumX = 0, sumY = 0, sumXX = 0, sumXY = 0;

    for (size_t i = 0; i < N; ++i) {
        sumX += x[i];
        sumY += y[i];
        sumXX += x[i] * x[i];
        sumXY += x[i] * y[i];
    }

    // Compute the coefficients using the normal equation:
    // a0 = (sumY * sumXX - sumX * sumXY) / (N * sumXX - sumX * sumX)
    // a1 = (N * sumXY - sumX * sumY) / (N * sumXX - sumX * sumX)

    TYPE denominator = N * sumXX - sumX * sumX;
    if (denominator == 0) {
        throw std::runtime_error("Denominator in normal equation is zero. This indicates a degenerate case.");
    }

    TYPE a0 = (sumY * sumXX - sumX * sumXY) / denominator;
    TYPE a1 = (N * sumXY - sumX * sumY) / denominator;

    // Store the coefficients in the output vector
    coeffs.clear();
    coeffs.push_back(a0);
    coeffs.push_back(a1);

    return true;
}
////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
